/*
 * AIEngine a new generation network intrusion detection system.
 *
 * Copyright (C) 2013-2023  Luis Campo Giralte
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Library General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Library General Public License for more details.
 *
 * You should have received a copy of the GNU Library General Public
 * License along with this library; if not, write to the
 * Free Software Foundation, Inc., 51 Franklin St, Fifth Floor,
 * Boston, MA  02110-1301, USA.
 *
 * Written by Luis Campo Giralte <luis.camp0.2009@gmail.com>
 *
 */
#include "StackLanIPv6.h"

namespace aiengine {

StackLanIPv6::StackLanIPv6() {

	name("Lan IPv6 network stack");

	// Add the specific Protocol objects
	addProtocol(eth, true);
	addProtocol(vlan, false);
	addProtocol(mpls, false);
	addProtocol(pppoe, false);
	addProtocol(ip6_, true);
	addProtocol(tcp_, true);
	addProtocol(udp_, true);
	addProtocol(icmp6_, true);

	// Layer7 protocols
        addProtocol(http);
        addProtocol(ssl);
        addProtocol(smtp);
        addProtocol(imap);
        addProtocol(pop);
        addProtocol(ssh);
        addProtocol(bitcoin);
        addProtocol(mqtt);
        addProtocol(tcp_generic);
        addProtocol(freqs_tcp, false);
        addProtocol(dns);
        addProtocol(sip);
        addProtocol(ntp);
        addProtocol(snmp);
        addProtocol(ssdp);
	addProtocol(netbios);
        addProtocol(coap);
        addProtocol(rtp);
        addProtocol(quic);
        addProtocol(dtls);
        addProtocol(dhcp6_);
        addProtocol(udp_generic);
        addProtocol(freqs_udp, false);

        // Link the FlowCaches to their corresponding FlowManager for timeouts
        flow_table_udp_->setFlowCache(flow_cache_udp_);
        flow_table_tcp_->setFlowCache(flow_cache_tcp_);

	// Link the protocols with the multiplexers
	ip6_->setMultiplexer(mux_ip);
	mux_ip->setProtocol(static_cast<ProtocolPtr>(ip6_));

	icmp6_->setMultiplexer(mux_icmp_);
	mux_icmp_->setProtocol(static_cast<ProtocolPtr>(icmp6_));

	udp_->setMultiplexer(mux_udp_);
	mux_udp_->setProtocol(static_cast<ProtocolPtr>(udp_));
	ff_udp_->setProtocol(static_cast<ProtocolPtr>(udp_));

	tcp_->setMultiplexer(mux_tcp_);
	mux_tcp_->setProtocol(static_cast<ProtocolPtr>(tcp_));
	ff_tcp_->setProtocol(static_cast<ProtocolPtr>(tcp_));

        dhcp6_->setFlowForwarder(ff_dhcp6_);
        ff_dhcp6_->setProtocol(static_cast<ProtocolPtr>(dhcp6_));

	// Configure the multiplexers
	mux_eth->addUpMultiplexer(mux_ip);
	mux_ip->addDownMultiplexer(mux_eth);
	mux_ip->addUpMultiplexer(mux_udp_);
	mux_udp_->addDownMultiplexer(mux_ip);
	mux_ip->addUpMultiplexer(mux_tcp_);
	mux_tcp_->addDownMultiplexer(mux_ip);
	mux_ip->addUpMultiplexer(mux_icmp_);
	mux_icmp_->addDownMultiplexer(mux_ip);

	// Connect the FlowManager and FlowCache
	tcp_->setFlowCache(flow_cache_tcp_);
	tcp_->setFlowManager(flow_table_tcp_);
	flow_table_tcp_->setProtocol(tcp_);

	udp_->setFlowCache(flow_cache_udp_);
	udp_->setFlowManager(flow_table_udp_);
	flow_table_udp_->setProtocol(udp_);

        // Connect to upper layers the FlowManager
        http->setFlowManager(flow_table_tcp_);
        ssl->setFlowManager(flow_table_tcp_);
        smtp->setFlowManager(flow_table_tcp_);
        imap->setFlowManager(flow_table_tcp_);
        pop->setFlowManager(flow_table_tcp_);
        mqtt->setFlowManager(flow_table_tcp_);
        ssh->setFlowManager(flow_table_tcp_);

        dns->setFlowManager(flow_table_udp_);
        sip->setFlowManager(flow_table_udp_);
        ssdp->setFlowManager(flow_table_udp_);
        coap->setFlowManager(flow_table_udp_);
        netbios->setFlowManager(flow_table_udp_);
        dhcp6_->setFlowManager(flow_table_udp_);
        dtls->setFlowManager(flow_table_udp_);
        quic->setFlowManager(flow_table_udp_);

        freqs_tcp->setFlowManager(flow_table_tcp_);
        freqs_udp->setFlowManager(flow_table_udp_);

	// Connect the AnomalyManager with the protocols that may have anomalies
        ip6_->setAnomalyManager(anomaly_);
        tcp_->setAnomalyManager(anomaly_);
        udp_->setAnomalyManager(anomaly_);

	// Configure the FlowForwarders
	tcp_->setFlowForwarder(ff_tcp_);
	udp_->setFlowForwarder(ff_udp_);

        setTCPDefaultForwarder(ff_tcp_);
        setUDPDefaultForwarder(ff_udp_);

        std::ostringstream msg;
        msg << name() << " ready.";

        infoMessage(msg.str());

	setMode("full");
}

void StackLanIPv6::showFlows(std::basic_ostream<char> &out, std::function<bool (const Flow&)> condition, int protocol) const {

        int total = flow_table_tcp_->getTotalFlows() + flow_table_udp_->getTotalFlows();
        out << "Flows on memory " << total << std::endl;
        if (protocol == IPPROTO_TCP)
                flow_table_tcp_->showFlows(out, condition);
        else if (protocol == IPPROTO_UDP)
                flow_table_udp_->showFlows(out, condition);
        else {
                flow_table_tcp_->showFlows(out, condition);
                flow_table_udp_->showFlows(out, condition);
        }
}

void StackLanIPv6::showFlows(Json &out, std::function<bool (const Flow&)> condition, int protocol) const {

        if (protocol == IPPROTO_TCP)
                flow_table_tcp_->showFlows(out, condition);
        else if (protocol == IPPROTO_UDP)
                flow_table_udp_->showFlows(out, condition);
        else {
                flow_table_tcp_->showFlows(out, condition);
                flow_table_udp_->showFlows(out, condition);
	}
}

void StackLanIPv6::setMode(const std::string &mode) {

        std::initializer_list<SharedPointer<FlowForwarder>> tcp_list {
                ff_tcp_, ff_http, ff_ssl, ff_smtp, ff_imap, ff_pop, ff_bitcoin,
                ff_mqtt, ff_ssh, ff_tcp_generic, ff_tcp_freqs};
        std::initializer_list<SharedPointer<FlowForwarder>> udp_list {
                ff_udp_, ff_dns, ff_sip, ff_ntp, ff_snmp, ff_ssdp, ff_netbios,
                ff_coap, ff_rtp, ff_quic, ff_dtls, ff_dhcp6_, ff_udp_generic, ff_udp_freqs};

        if (auto it = modes.find(mode); it != modes.end()) {
                std::ostringstream msg;

                disableFlowForwarders(tcp_list);
                disableFlowForwarders(udp_list);

                if ((*it).second == Mode::FULL) {

			operation_mode_ = "full";
                        msg << "Enable FullEngine mode on " << name();
                        enableFlowForwarders(tcp_list);
                        enableFlowForwarders(udp_list);

                } else if ((*it).second == Mode::FREQUENCIES) {

			operation_mode_ = "frequency";
                        msg << "Enable FrequencyEngine mode on " << name();
                        enableFlowForwarders({ff_tcp_, ff_tcp_freqs});
                        enableFlowForwarders({ff_udp_, ff_udp_freqs});

                } else if ((*it).second == Mode::NIDS) {

			operation_mode_ = "nids";
                        msg << "Enable NIDSEngine mode on " << name();
                        enableFlowForwarders({ff_tcp_, ff_tcp_generic});
                        enableFlowForwarders({ff_udp_, ff_udp_generic});
                }
                infoMessage(msg.str());
        }
}

void StackLanIPv6::setTotalTCPFlows(int value) {

	flow_cache_tcp_->create(value);
	tcp_->createTCPInfos(value);

	// The vast majority of the traffic of internet is HTTP
	// so create 75% of the value received for the http caches
	http->increaseAllocatedMemory(value * 0.75);

        // The 40% of the traffic is SSL
        ssl->increaseAllocatedMemory(value * 0.4);

	// 5% of the traffic could be SMTP/IMAP, im really positive :D
	smtp->increaseAllocatedMemory(value * 0.05);
	imap->increaseAllocatedMemory(value * 0.05);
	pop->increaseAllocatedMemory(value * 0.05);
	bitcoin->increaseAllocatedMemory(value * 0.05);
	mqtt->increaseAllocatedMemory(value * 0.05);
	ssh->increaseAllocatedMemory(value * 0.05);
	freqs_tcp->increaseAllocatedMemory(16);
}

void StackLanIPv6::setTotalUDPFlows(int value) {

	flow_cache_udp_->create(value);

	dns->increaseAllocatedMemory(value/ 2);
        sip->increaseAllocatedMemory(value * 0.2);
        ssdp->increaseAllocatedMemory(value * 0.2);
        coap->increaseAllocatedMemory(value * 0.2);
        dtls->increaseAllocatedMemory(value * 0.2);
        quic->increaseAllocatedMemory(value * 0.2);
        dhcp6_->increaseAllocatedMemory(value * 0.1);
	freqs_udp->increaseAllocatedMemory(16);
}

int StackLanIPv6::getTotalTCPFlows() const { return flow_cache_tcp_->getTotalFlows(); }

int StackLanIPv6::getTotalUDPFlows() const { return flow_cache_udp_->getTotalFlows(); }

void StackLanIPv6::setFlowsTimeout(int timeout) {

        flow_table_udp_->setTimeout(timeout);
        flow_table_tcp_->setTimeout(timeout);
}

void StackLanIPv6::setTCPRegexManager(const SharedPointer<RegexManager> &rm) {

        tcp_->setRegexManager(rm);
        tcp_generic->setRegexManager(rm);
	super_::setTCPRegexManager(rm);
}

void StackLanIPv6::setUDPRegexManager(const SharedPointer<RegexManager> &rm) {

        udp_->setRegexManager(rm);
        udp_generic->setRegexManager(rm);
	super_::setUDPRegexManager(rm);
}

void StackLanIPv6::setTCPIPSetManager(const SharedPointer<IPSetManager> &ipset_mng) {

	tcp_->setIPSetManager(ipset_mng);
	super_::setTCPIPSetManager(ipset_mng);
}

void StackLanIPv6::setUDPIPSetManager(const SharedPointer<IPSetManager> &ipset_mng) {

	udp_->setIPSetManager(ipset_mng);
	super_::setUDPIPSetManager(ipset_mng);
}

void StackLanIPv6::stopAsioService() {

#if defined(HAVE_REJECT_FLOW)
	rj_mng_->close();
#endif
}

void StackLanIPv6::setAsioService(boost::asio::io_service &io_service, const std::string &device) {

        // Create a new RejectManager with their corresponding sockets
#if defined(HAVE_REJECT_FLOW)
        if (geteuid() == 0) { // The process have rights on raw sockets
		rj_mng_->setAsioService(io_service, device);
                if (rj_mng_->ready()) {
                        // Attach the reject function to the corresponding protocols tcp/udp
                        tcp_->addRejectFunction(std::bind(&RejectManager<StackLanIPv6>::rejectTCPFlow, rj_mng_, std::placeholders::_1));
                        udp_->addRejectFunction(std::bind(&RejectManager<StackLanIPv6>::rejectUDPFlow, rj_mng_, std::placeholders::_1));
                }
        }
#endif

}

void StackLanIPv6::statistics(std::basic_ostream<char> &out) const {

        super_::statistics(out);
#if defined(HAVE_REJECT_FLOW)
        if (geteuid() == 0) { // The process have rights on raw sockets
                if (rj_mng_) {
                        rj_mng_->statistics(out);
                        out << std::endl;
                }
        }
#endif
}

#if defined(JAVA_BINDING)

void StackLanIPv6::setTCPRegexManager(RegexManager *sig) {

        SharedPointer<RegexManager> rm;

        if (sig != nullptr)
                rm.reset(sig);

        setTCPRegexManager(rm);
}

void StackLanIPv6::setUDPRegexManager(RegexManager *sig) {

        SharedPointer<RegexManager> rm;

        if (sig != nullptr)
                rm.reset(sig);

        setUDPRegexManager(rm);
}

void StackLanIPv6::setTCPIPSetManager(IPSetManager *ipset_mng) {

        SharedPointer<IPSetManager> im;

        if (ipset_mng != nullptr)
                im.reset(ipset_mng);

        setTCPIPSetManager(im);
}

void StackLanIPv6::setUDPIPSetManager(IPSetManager *ipset_mng) {

        SharedPointer<IPSetManager> im;

        if (ipset_mng != nullptr)
                im.reset(ipset_mng);

        setUDPIPSetManager(im);
}

#endif

std::tuple<Flow*, Flow*> StackLanIPv6::getCurrentFlows() const {

        Flow *flow = nullptr;
	uint16_t proto = ip6_->getL7Protocol();

        if (proto == IPPROTO_TCP)
                flow = tcp_->getCurrentFlow();
        else if (proto == IPPROTO_UDP)
                flow = udp_->getCurrentFlow();

#if GCC_VERSION < 50500
        return std::tuple<Flow*, Flow*>(flow, nullptr);
#else
        return {flow, nullptr};
#endif
}

SharedPointer<Flow> StackLanIPv6::getFlow(const FlowSearchOptions &fsos) const {

        if (fsos.protocol == IPPROTO_UDP)
                return flow_table_udp_->find(fsos);
        else
                return flow_table_tcp_->find(fsos);
}

} // namespace aiengine
